
import os
import pprint
from abc import abstractmethod
from typing import List, Dict, Tuple

# 标准化 还不成熟
from .utils import *
from .constants import *

class Config_template():
    def __init__(self, config_path:str, config_base_name:str='Configbase'):
        self.config_path = config_path
        self.config_base_name = config_base_name
        if os.path.exists(config_path):
            self.configbase = read_config(config_path)
        else:
            self.configbase = {}
            raise ValueError("Error: config_file not found")
    @abstractmethod
    def get_config(self, config_id)->Dict:
        raise NotImplementedError()
    def get_config_by_type(self,config_type:str, config_name)->Dict:
        # if 'id' in config_type:
        #     return self.configbase[basename][config_name]
        return next((item for item in self.configbase[self.config_base_name] if item[config_type] == config_name), None)

    def get_configs_by_type(self,config_type:str, config_name)->List[Dict]:
        return [item for item in self.configbase[self.config_base_name] if item[config_type] == config_name]


class Dataset(Config_template):
    def __init__(self, config_path:str, dataset_id:int=0, dataset_name:str=None, dataset_dict:Dict=None):
        super().__init__(config_path, 'Datasetbase')
        
        if dataset_id is None and dataset_name is None: 
            self.dataset_id = dataset_dict.get('dataset_id',dataset_id)
            self.dataset_name = dataset_dict.get('dataset_name',dataset_name)
            self.dataset = dataset_dict
        elif dataset_name is None:
            self.dataset_id = dataset_id
            self.dataset = self.get_config(dataset_id)
            self.dataset_name = self.dataset.get('dataset_name',dataset_name)
        else: # dataset_id is None and dataset_name is not None or both is not None
            self.dataset_name = dataset_name
            self.dataset = self.get_config_by_name(dataset_name)
            self.dataset_id = self.dataset.get('dataset_id',dataset_id)

    def get_config(self, dataset_id:int)->Dict:
        return self.get_config_by_type('dataset_id', dataset_id)

    def get_config_by_name(self, dataset_name)->Dict:
        return self.get_config_by_type('dataset_name', dataset_name)

    def get_dataset_description(self)->str:
        dataset_description = self.dataset.get('description',None)
        dataset_description_path = self.dataset.get('description_path',None)
        if isempty(dataset_description) and not isempty(dataset_description_path):
            origin = read_txt(dataset_description_path)
            return origin
        elif isempty(dataset_description_path) and not isempty(dataset_description):
            return dataset_description
        else:
            print("Error: dataset_description and dataset_description_path cannot be both empty")
            return ""


class Space(Config_template):
    def __init__(self, config_path:str, space_id:int=0, space_dict:Dict=None):
        super().__init__(config_path, 'Spacebase')
        self.space_id = space_id
        if space_id is None: 
            self.space = space_dict
            self.input_vars = []
            self.space_description = ""
        else:
            self.space = self.get_config(space_id)
            self.input_vars = self.get_total_invars()
            self.space_description = self.get_space_description()

    def get_config(self, space_id:int)->Dict:
        return self.get_config_by_type('space_id', space_id)

    def get_total_invars(self)->List[str]:
        return self.space.get('input_vars',[])

    def get_complete_description(self,space_dict:Dict=None)->str:
        """ space_description中 不包含 字典{}时可用

        Args:
            space_dict (Dict, optional): _description_. Defaults to None.

        Returns:
            str: _description_
        """
        
        assert not isempty(self.space_description), "Error: space_description cannot be empty"
        
        if isempty(self.input_vars):
            return self.space_description
        else:
            list_set = set(self.input_vars)
            config_keys_set = set(space_dict.keys())
            all_in_keys = list_set.issubset(config_keys_set)
            if all_in_keys:
                return self.space_description.format(**space_dict)
            else:
                print("Error: input_vars is not empty, config lost some keys")
                return ""
    def get_space_description(self)->str:
        space_description = self.space.get('description',None)
        space_description_relevant_path = self.space.get('description_path',None)
        space_description_path = os.path.join(os.path.dirname(space_path), space_description_relevant_path) if not isempty(space_description_relevant_path) else ""
        if isempty(space_description) and os.path.exists(space_description_path):
            origin = read_txt(space_description_path)
            result = split_text_by_catagory(origin, '-->',suffix=True,contain=False)
            return result
        elif isempty(space_description_path) and not isempty(space_description):
            return space_description
        else:
            print("Error: space_description and space_description_path cannot beboth empty")
            return ""
        
    
    def retrive_knowledge_by_space(self) -> List[str]:
        space_id = self.space['space_id']
        knowledgebase = Knowledge(knowledge_path)
        knowledge_list = knowledgebase.retrive_knowledges_by_space(space_id)
        return knowledge_list
    
    def generate_request_suffix(self)->str:
        assert 'output_schema' in self.space 
        output_dict = self.space['output_schema']

        assert 'input_vars' in output_dict and 'request_cmd' in output_dict
        output_invars = output_dict['input_vars']
        request_str = output_dict['request_cmd']
        # 需要配置
        if not isempty(output_invars): 
            config_dict = {k:output_dict[k] for k in output_invars}
            return request_str.format(**config_dict)
        else:
            return request_str
    
    def get_example_str(self)->str:
        assert 'output_schema' in self.space 
        output_dict = self.space['output_schema']
        # assert 'input_vars' in output_dict
        # output_invars = output_dict['input_vars']
        assert 'examples' in output_dict
        return output_dict['examples']
        

class Prompt(Config_template):
    def __init__(self, config_path:str, prompt_id:int=0):
        super().__init__(config_path, 'Promptbase')
        self.prompt_id = prompt_id
        self.prompt_dict = self.get_config(prompt_id)

    def get_config(self, prompt_id:int)->Dict:
        return self.get_config_by_type('prompt_id', prompt_id)

    def get_config_by_name(self, prompt_name:str)->Dict:
        return self.get_config_by_type('prompt_name', prompt_name)
    
    def get_prompt_template_str(self):
        prompt_template_str = ''
        if self.prompt_dict is not None:
            prompt_template = self.prompt_dict.get('prompt_template',None)
            prompt_template_relevant_path = self.prompt_dict.get('prompt_template_path',None)
            prompt_template_path = os.path.join(os.path.dirname(prompt_path), prompt_template_relevant_path) if not isempty(prompt_template_relevant_path) else ""
            
            if isempty(prompt_template) and os.path.exists(prompt_template_path):
                origin = read_txt(prompt_template_path)
                prompt_template_str = self.refine_prompt(origin)
            elif isempty(prompt_template_path) and not isempty(prompt_template):
                prompt_template_str = self.refine_prompt(prompt_template)
            else:
                print("Error: prompt_template and prompt_template_path cannot be both empty")
        return prompt_template_str

    def get_prompt_dict(self) -> Dict:
        return self.prompt_dict
    def get_prompt_item(self, key):
        return self.prompt_dict.get(key,None)

    def get_suffix_prefix_by_name(self, catagory) -> Tuple[str, str]:
        prompt_str = self.get_prompt_template_str()
        return split_text_by_catagory(prompt_str, catagory, suffix=False, contain=False), \
               split_text_by_catagory(prompt_str, catagory, suffix=True, contain=False)
      
    @staticmethod
    def refine_prompt(origin:str)->str:
        # 使用正则表达式提取 #Role  后面(且包含）的内容
        catagory = r'# Role'
        return split_text_by_catagory(origin, catagory, suffix=True, contain=True)
    
class Knowledge(Config_template):
    def __init__(self, config_path:str):
        super().__init__(config_path, 'Knowledgebase')

    def get_config(self, knowledge_id:int)->Dict:
        return self.get_config_by_type('knowledge_id', knowledge_id)
    
    def get_config_by_space(self,space_id:int=0)->List[Dict]:
        result = []
        for item in self.configbase[self.config_base_name]:
            space_id_temp = item['space_id']
            if isinstance(space_id_temp, int) and space_id_temp == space_id:
                result.append(item)
            elif isinstance(space_id_temp, str) and space_id_temp.strip() == str(space_id):
                result.append(item)
            elif isinstance(space_id_temp, list) and space_id in space_id_temp:
                result.append(item)
        return result
    
    def retrive_knowledges_by_space(self, space_id:int=None)->List[str]:
        assert space_id is not None, "space_id is None, not support to get_knowledges_by_space"
        space_knowledges = self.get_config_by_space(space_id)
        return self.combine_knowledge(space_knowledges)

    @staticmethod
    def combine_knowledge(knowledge_list:List[Dict])->List[str]:
        return [item['description'] for item in knowledge_list]

class Constraints:
    def __init__(self, contraints:Dict=None):
        if contraints is None: 
            self.input_vars = []
            self.description = ""
        else:
            self.description = contraints['description']
            self.input_vars = contraints['input_vars']

    def get_complete_str(self, config:Dict=None)->str:
        if len(self.input_vars)==0:
            return self.description
        elif config is None:
            print("Error: input_vars is not empty, but config is empty")
        else: 
            list_set = set(self.input_vars)
            config_keys_set = set(config.keys())
            all_in_keys = list_set.issubset(config_keys_set)
            if all_in_keys:
                return self.description.format(**config)
            else:
                print("Error: input_vars is not empty, config lost some keys")
        return None

class Question:
    def __init__(self, question:Dict=None):
        if question is None: 
            self.request = ""
        else:
            self.request = question.get('request',"")
    
    def get_request(self)->str:
        if isempty(self.request): 
            return ""
        return self.request
    pass

class Evallog:
    template = "-epoch: {id},\n-configurations: {config},\n-metrics: {metrics}\n"
    input_vars = ["id", "config", "metrics"]
    def __init__(self, evallog:List[str]=None):
        if evallog is None:
            self.evallog = []
        else:
            self.evallog = evallog        
    def generate_log(self, id:int, config:Dict, metrics:Dict)->str:
        config_str = represent_dict(config)
        metrics_str = represent_dict(metrics)
        constraints_dict = {"id":id, "config":config_str, "metrics":metrics_str}
        return self.template.format(**constraints_dict)
    
    def add_log(self, id:int, config:Dict, metrics:Dict)->None:
        self.evallog.append(self.generate_log(id, config, metrics))
    
    def clean_log(self)->None:
        self.evallog = []
    
    def __len__(self)->int:
        return len(self.evallog)
    
    def __repr__(self):
        return "\n".join(self.evallog)
    
    def isempty(self)->bool:
        return len(self.evallog) == 0


def represent_dict(d: Dict) -> str:
    formatted_str_list = [f"'{k}': {v}" for k, v in d.items()]
    return ", ".join(formatted_str_list)
    # formatted_str = pprint.pformat(d)
    # # 去除首尾的 {}
    # if formatted_str.startswith('{') and formatted_str.endswith('}'):
    #     formatted_str = formatted_str[1:-1]
    # return formatted_str

class Task(Config_template):
    def __init__(self, config_path:str, task_id:int=0, task_dict:Dict=None):
        super().__init__(config_path, 'Taskbase')

        self.task_id = task_id
        if task_id is None: 
            self.task = task_dict
        else:
            self.task = self.get_config(task_id)

        if self.task is None:
            self.usable = False
            self.question = None
            self.dataset = None
            self.constraints = None
            self.knowledge = None
            self.space = None
            self.evallog = None
            self.fix_len = 0
        else:
            self.question = self.get_question()
            self.evallog = self.get_evallog()
            self.space = self.get_space()
            self.dataset = self.get_dataset()
            self.constraints = self.get_constraints()
            self.knowledge, self.fix_len = self.get_knowledge(self.space)

            self.usable = True
    def get_config(self, task_id:int)->Dict:
        return self.get_config_by_type('task_id', task_id)

    def get_question(self)->Question:
        question = self.task['Question']
        return Question(question)
    
    def get_space(self)->Space:
        space_dict = self.task['Space']
        if space_dict is not None:
            space_id = space_dict['space_id']
            if eval(space_dict['indatabase']):
                return Space(space_path,space_id)
        return Space(space_path,None,space_dict)

    def get_dataset(self)->Dataset:
        dataset_dict = self.task.get('Dataset',None)
        if dataset_dict is not None:
            dataset_id = dataset_dict.get('dataset_id',None)
            dataset_name = dataset_dict.get('dataset_name',None)
            dataset_id = None if isempty(dataset_id) else dataset_id
            dataset_name = None if isempty(dataset_name) else dataset_name
            if eval(dataset_dict['indatabase']):            
                return Dataset(dataset_path, dataset_id, dataset_name, dataset_dict)
        return Dataset(dataset_path, None, None, dataset_dict)
    
    def get_knowledge(self, space:Space) -> Tuple[List[str],int]:
        # retrieve knowledge from space by space_id
        knowledge_list = space.retrive_knowledge_by_space()
        dynamic_len = len(knowledge_list)
        # add new knowledge into knowledge_list
        knowledge_extra_dict = self.task['Knowledge']
        if not isempty(knowledge_extra_dict):
            assert isinstance(knowledge_extra_dict, dict)
            knowledge_extra_list = knowledge_extra_dict.get('content',None)
            assert not isempty(knowledge_extra_list)
            if isinstance(knowledge_extra_list, list):
                knowledge_list.extend(knowledge_extra_list)
            elif isinstance(knowledge_extra_list, str):
                knowledge_list.append(knowledge_extra_list)
            else:
                raise ValueError("Error: knowledge_extra_list is not a list or a string")
        return knowledge_list,dynamic_len
    def get_constraints(self)->Constraints:
        constraints = self.task['Constraints']
        if isempty(constraints):
            return None
        return Constraints(constraints)

    def get_evallog(self)->Evallog:
        evallog = self.task['EvalLogs']
        if isempty(evallog):
            return Evallog()
        return Evallog(evallog)

    def generate_requirement(self)->str:
        requirement = self.question.get_request() + "\n"
        requirement += self.space.generate_request_suffix()
        return requirement
    
